package saramax

import (
	logger2 "booksmart/book/pkg/logger"
	"context"
	"encoding/json"
	"github.com/IBM/sarama"
	"time"
)

type BatchConsumer[T any] struct {
	l             logger2.LoggerV1
	fn            func(msgs []*sarama.ConsumerMessage, ts []T) error
	batchSize     int
	batchDuration time.Duration
}

func NewBatchConsumer[T any](l logger2.LoggerV1, fn func(msgs []*sarama.ConsumerMessage, ts []T) error) *BatchConsumer[T] {
	return &BatchConsumer[T]{
		l:             l,
		fn:            fn,
		batchSize:     10,
		batchDuration: 10 * time.Second,
	}
}

func (b *BatchConsumer[T]) Setup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (b *BatchConsumer[T]) Cleanup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (b *BatchConsumer[T]) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	msgCh := claim.Messages()
	batchSize := b.batchSize
	for {
		ctx, cancel := context.WithTimeout(context.Background(), b.batchDuration)
		var last *sarama.ConsumerMessage
		done := false
		msgs := make([]*sarama.ConsumerMessage, 0, batchSize)
		ts := make([]T, 0, batchSize)
		for i := 0; i < batchSize && !done; i++ {
			select {
			case <-ctx.Done():
				done = true
			case msg := <-msgCh:
				last = msg
				var t T
				err := json.Unmarshal(msg.Value, &t)
				if err != nil {
					b.l.Error("反序列化失败", logger2.Int64("partition", msg.Partition),
						logger2.String("topic", msg.Topic),
						logger2.Int64("offset", msg.Offset),
						logger2.Error(err))
					continue
				}
				msgs = append(msgs, msg)
				ts = append(ts, t)
			}
		}
		cancel()
		if len(ts) == 0 {
			continue
		}
		err := b.fn(msgs, ts)
		if err != nil {
			b.l.Error("批量消费失败", logger2.Error(err))
		}

		if last != nil {
			session.MarkMessage(last, "")
		}

	}

}
